# encoding: UTF-8
"""将不同时间框架（不同周期）下的价格走势分为：升、盘、降（XPJ)，或者涨、调、跌（UAD）
低频下的分类决定交易方向，高频下的决定进场时机
（多、币、空）=（持多仓、持币、持空仓）=(B, N, S)=(L, M, S)
B=BUY, N=NULL, S=SELL
L=LONG, M=MONEY, S=SHORT
UAD=(UP,ADJUST,DOWN)
XPJ=(XING,PAN, JIANG)
"""
import numpy as np
import pandas as pd

from base.backtest_result import calc_long_trade_result, calc_short_trade_result
from strategy.istrategy import IStrategy
from const_struct.constant import *
from indicator.indicators import AMA


class TrendStrategy(IStrategy):
    # 参数缺省值
    p1 = 5
    p2 = 24
    p3 = 120
    p4 = 600
    class_name = 'TrendStrategy'

    def __init__(self, strategy_event_engine, setting):
        super().__init__(strategy_event_engine, setting)

    # def on_init(self, params=None):
    #     super().on_init(params=params)
    #     if not params and not self.params:
    #         self.load_params()
    #     if self.period == PERIOD_TYPE_TICK:
    #         self.load_tick(days=10)
    #     else:
    #         self.load_bar(days=60)
    #     print('TrendStrategy on_init', type(self.ticks), self.ticks.__dict__)

    def on_tick(self, tick):
        """
        run by strategy_engine, according to new tick,
        calculate the variables
        :param tick:
        :return:
        """
        print(tick.__dict__)

    def load_params(self):
        super().load_params()
        # print(self.params)

    def optimise_params(self):
        super().optimise_params()

    def on_order(self, order):
        pass

    def on_trade(self, trade):
        pass

    def on_bar(self, bar):
        pass

    def on_bars(self, bars,**kwargs):
        params = {'p1':self.p1,'p2':self.p2,'p3':self.p3,'p4':self.p4}
        self.out_data=self.calc_long_short_signal(bars,params=params)

    @classmethod
    def calc_long_signal(cls, data, params):
        return None

    @classmethod
    def calc_short_signal(cls, data, params):
        return None

    @classmethod
    def calc_long_short_signal(cls, data, params, period=PERIOD_TYPE_TICK, buy_cond=1, sell_cond=2, short_cond=3,
                               cover_cond=4):
        try:
            p1 = params['p1']
            p2 = params['p2']
            p3 = params['p3']
            p4 = params['p4']

        except KeyError:
            return {'error': '参数错误'}
        print('calc_long_short_signal:', data.columns)
        if period == PERIOD_TYPE_TICK:
            tprice = np.log(data['lastPrice'])
            sgn_close = 'lastPrice'
        else:
            tprice = np.log(data['close'])
            sgn_close = 'close'
        # print('calc_long_short_signal:', tprice.head())
        data['logp'] = tprice
        out_data = data
        log_price = data['logp']
        profit = log_price.diff(1)  # 1周期的持多利润率（收益率）
        profit[0] = 0
        data_len = len(profit)
        # print('calc_long_short_signal:', data.head(100))
        # 第一笔开多1手后的持有利润率
        total_long_profit = profit.rolling(window=data_len, min_periods=1).sum()
        ama = total_long_profit  # AMA(total_long_profit)
        # 持有 p4周期的利润率
        # sum_profit = profit.rolling(window=p4, min_periods=1).sum()
        # p4周期内的最高利润率
        mx = ama.rolling(window=p4, min_periods=1).max()
        # p4周期内的最低利润率
        mn = ama.rolling(window=p4, min_periods=1).min()
        # p周期的平均利润率
        m1 = ama.rolling(window=p1, min_periods=1).mean()
        m2 = ama.rolling(window=p2, min_periods=1).mean()
        m3 = ama.rolling(window=p3, min_periods=1).mean()
        m4 = ama.rolling(window=p4, min_periods=1).mean()
        dama = ama.diff(1)
        dmx = mx.diff(1)
        dmn = mn.diff(1)
        dm1 = m1.diff(1)
        dm2 = m2.diff(1)
        dm3 = m3.diff(1)
        dm4 = m4.diff(1)
        price1 = tprice[dmx > 0]  #
        price2 = tprice[dmn < 0]  #
        c1 = pd.Series(data=[1] * len(price1), index=price1.index)  #
        c2 = pd.Series(data=[1] * len(price2), index=price2.index)  #
        price3 = tprice[dmx < 0]  #
        price4 = tprice[dmn > 0]  #
        c3 = pd.Series(data=[1] * len(price3), index=price3.index)  #
        c4 = pd.Series(data=[1] * len(price4), index=price4.index)  #
        buy_series = c1
        sell_series = c2
        short_series = c2
        cover_series = c1
        if buy_cond == 1:
            buy_series = c1
        elif buy_cond == 2:
            buy_series = c2
        elif buy_cond == 3:
            buy_series = c3
        elif buy_cond == 4:
            buy_series = c4
        if sell_cond == 1:
            sell_series = c1
        elif sell_cond == 2:
            sell_series = c2
        elif sell_cond == 3:
            sell_series = c3
        elif sell_cond == 4:
            sell_series = c4
        if short_cond == 1:
            short_series = c1
        elif short_cond == 2:
            short_series = c2
        elif short_cond == 3:
            short_series = c3
        elif short_cond == 4:
            short_series = c4
        if cover_cond == 1:
            cover_series = c1
        elif cover_cond == 2:
            cover_series = c2
        elif cover_cond == 3:
            cover_series = c3
        elif cover_cond == 4:
            cover_series = c4

        out_data['buy'] = buy_series
        out_data['sell'] = sell_series
        out_data['short'] = short_series
        out_data['cover'] = cover_series
        out_data['tprice'] = data[sgn_close]
        # print('long_short_signal ok!',out_data.query('buy==1 or sell==1'))
        print(out_data.head())
        return out_data

    @classmethod
    def calc_trade_result(cls, data, params, long_short=1, recalc_signal=0, period=PERIOD_TYPE_TICK):
        """
        计算分析交易（策略）结果
        :param data:
        :param params: 计算统计结果要用到的参数（见calc_long_trade_result的说明 ），
        当recalc_signal==1时，同时包含策略参数
        :param long_short:
        :param recalc_signal: 0=不重新计算交易信号，
        :return:
        """
        # print('TrendStrategy.calc_trade_result', data.head(), params, long_short == 1, recalc_signal == 0)
        rslt = {'long': {}, 'short': {}}
        sgn_close = 'lastPrice' if period == PERIOD_TYPE_TICK else 'close'
        if recalc_signal != 0:
            if long_short == 0:
                data = cls.calc_long_short_signal(data, params, period=period)
                # print(data.head(),data.columns,sgn_close)
                rslt['long'] = calc_long_trade_result(data, sgn_close_price=sgn_close)
                rslt['short'] = calc_short_trade_result(data, sgn_close_price=sgn_close)
                return rslt
            elif long_short == 1:
                data = cls.calc_long_signal(data, params)
                rslt['long'] = calc_long_trade_result(data, sgn_close_price=sgn_close)
            else:
                data = cls.calc_short_signal(data, params)
                rslt['short'] = calc_short_trade_result(data, sgn_close_price=sgn_close)
        else:
            if long_short == 0:
                rslt['long'] = calc_long_trade_result(data, sgn_close_price=sgn_close)
                rslt['short'] = calc_short_trade_result(data, sgn_close_price=sgn_close)
                return rslt
            elif long_short == 1:
                # print(data.head())
                rslt['long'] = calc_long_trade_result(data, sgn_close_price=sgn_close)
            else:
                rslt['short'] = calc_short_trade_result(data, sgn_close_price=sgn_close)
        return rslt

    def run_param_optimization(cls, data, params):
        super().run_param_optimization(data, params)

    def on_ticks(self, ticks):
        """
        根据 ticks， 计算多空信号
        :param ticks: pd.DataFrame,包含：lastPrice，index:datetime
        :return:
        """
        p1 = self.p1
        p2 = self.p2
        p3 = self.p3
        p4 = self.p4
        ticks['logp'] = np.log(ticks['lastPrice'])
        price = ticks['lastPrice']
        self.out_data = pd.DataFrame(data=price, index=price.index, columns=['close'])
        log_price = ticks['logp']
        profit = log_price.diff(1)  # 1周期的持多利润率（收益率）
        profit[0] = 0
        data_len = len(profit)
        # print(data_len, price.head(), type(ticks))
        # 第一笔开多1手后的持有利润率
        total_long_profit = profit.rolling(window=data_len, min_periods=1).sum()
        ama = AMA(total_long_profit)
        # 持有 p4周期的利润率
        # sum_profit = profit.rolling(window=p4, min_periods=1).sum()
        # p4周期内的最高利润率
        mx = ama.rolling(window=p4, min_periods=1).max()
        # p4周期内的最低利润率
        mn = ama.rolling(window=p4, min_periods=1).min()
        # p周期的平均利润率
        m1 = ama.rolling(window=p1, min_periods=1).mean()
        m2 = ama.rolling(window=p2, min_periods=1).mean()
        m3 = ama.rolling(window=p3, min_periods=1).mean()
        m4 = ama.rolling(window=p4, min_periods=1).mean()
        dama = ama.diff(1)
        dmx = mx.diff(1)
        dmn = mn.diff(1)
        dm1 = m1.diff(1)
        dm2 = m2.diff(1)
        dm3 = m3.diff(1)
        dm4 = m4.diff(1)
        price1 = price[dmx > 0]  #
        price2 = price[dmn < 0]  #
        c1 = pd.Series(data=[1] * len(price1), index=price1.index)  #
        c2 = pd.Series(data=[1] * len(price2), index=price2.index)  #
        price3 = price[dmx < 0]  #
        price4 = price[dmn > 0]  #
        c3 = pd.Series(data=[1] * len(price3), index=price3.index)  #
        c4 = pd.Series(data=[1] * len(price4), index=price4.index)  #
        self.buy_series = c1
        self.sell_series = c2
        self.short_series = c2
        self.cover_series = c1
        if self.buy_cond == 1:
            self.buy_series = c1
        elif self.buy_cond == 2:
            self.buy_series = c2
        elif self.buy_cond == 3:
            self.buy_series = c3
        elif self.buy_cond == 4:
            self.buy_series = c4
        if self.sell_cond == 1:
            self.sell_series = c1
        elif self.sell_cond == 2:
            self.sell_series = c2
        elif self.sell_cond == 3:
            self.sell_series = c3
        elif self.sell_cond == 4:
            self.sell_series = c4
        if self.short_cond == 1:
            self.short_series = c1
        elif self.short_cond == 2:
            self.short_series = c2
        elif self.short_cond == 3:
            self.short_series = c3
        elif self.short_cond == 4:
            self.short_series = c4
        if self.cover_cond == 1:
            self.cover_series = c1
        elif self.cover_cond == 2:
            self.cover_series = c2
        elif self.cover_cond == 3:
            self.cover_series = c3
        elif self.cover_cond == 4:
            self.cover_series = c4

        self.out_data['buy'] = self.buy_series
        self.out_data['sell'] = self.sell_series
        self.out_data['short'] = self.short_series
        self.out_data['cover'] = self.cover_series


if __name__ == "__main__":
    from dbhelper.ctp_db_helper import *
    import matplotlib.pyplot as plt

    port = 27017
    host = '127.0.0.1'
    db_client = DbHelper(db_type=DbType.DBTYPE_MONGODB)

    ticks = pd.DataFrame()
    ticks = load_tick(db_client, 'rb1901', days=5, idx='datetime')  # , idx='datetime')
    prm = {'p1': 5, 'p2': 24, 'p3': 120, 'p4': 600}
    rslt = TrendStrategy.calc_trade_result(ticks, prm)
    print(rslt)
    # ticks['logp'] = np.log(ticks['lastPrice'])
    #
    # price = ticks['lastPrice']  # np.log(ticks['lastPrice'])
    #
    # # bars = load_bar(db_client, 'ni1901', days=10)
    # # bars['lnc'] = np.log(bars['close'])
    # #
    # # last_price = tick.query('volume>0')
    # p1 = 3
    # p2 = 15
    # p3 = 60
    # p4 = 120 * 15
    # # p1 = 5
    # # p2 = 24
    # # p3 = 120
    # # p4 = 2400
    # profit = price.diff(1)  # 前后1周期的持多利润率
    # profit[0] = 0
    # data_len = len(price)
    # # 第一笔开多1手后的持有利润率
    # total_long_profit = profit.rolling(window=data_len, min_periods=1).sum()
    # ama = AMA(total_long_profit, ksgn='lastPrice')
    # # 持有 p4周期的利润率
    # # sum_profit = profit.rolling(window=p4, min_periods=1).sum()
    # # 持有p4周期内的最高利润率
    # mx = ama.rolling(window=p4, min_periods=1).max()
    # # 持有p4周期内的最低利润率
    # mn = ama.rolling(window=p4, min_periods=1).min()
    # # 持有p周期的平均利润率
    # m1 = ama.rolling(window=p1, min_periods=1).mean()
    # m2 = ama.rolling(window=p2, min_periods=1).mean()
    # m3 = ama.rolling(window=p3, min_periods=1).mean()
    # m4 = ama.rolling(window=p4, min_periods=1).mean()
    # dama = ama.diff(1)
    # dmx = mx.diff(1)
    # dmn = mn.diff(1)
    # dmc = (dmx + dmn) / 2
    # dm1 = m1.diff(1)
    # dm2 = m2.diff(1)
    # dm3 = m3.diff(1)
    # dm4 = m4.diff(1)
    # print(total_long_profit.head(), ama.head())
    #
    # long = price[dmx > 0]
    # short = price[dmn < 0]
    # long_flag = pd.Series(data=[1] * len(long), index=long.index)
    # short_flag = pd.Series(data=[-1] * len(short), index=short.index)
    # long_short_flag = pd.concat([long_flag, short_flag])
    #
    # long = price[dmx > 0]  # 看多时的价格
    # short = price[dmn < 0]  # 看空时的价格
    # long_flag = pd.Series(data=[1] * len(long), index=long.index)  # 看多标志
    # short_flag = pd.Series(data=[-1] * len(short), index=short.index)  # 看空标志
    # bs = pd.concat([short_flag, long_flag])  # 多空标志
    # bs.sort_index(inplace=True)  # 重排序
    # sum_bs = bs.rolling(window=2, min_periods=1).sum()  # 若多空相随==0
    # sum_bs[0] = 0
    # bs = bs[sum_bs == 0]  # 去掉连续的多或空，仅保留第一次出现的
    # # strategy_profit = np.cumproduct(price[bs.index],bs)
    # # print('\n profit:\n',strategy_profit)
    #
    # buy = price[dmx > 0]
    # sell = price[dmn < 0]
    # buy_flag = pd.Series(data=[1] * len(buy), index=buy.index)
    # sell_flag = pd.Series(data=[-1] * len(sell), index=sell.index)
    # print(len(bs), bs.tail(), len(buy_flag), len(sell_flag))
    #
    # print(len(buy), buy.tail(), len(sell), sell.tail())
    # c1 = mx[dmx > 0]
    # c2 = mx[dmx < 0]
    # c3 = dmn[dmn > 0]
    # c4 = dmn[dmn < 0]
    #
    # # c5 = dmx[dmx > 0 and dmn > 0]
    # # c6 = dmx[dmx < 0 and dmn < 0]
    #
    # # ewm1 = tick['lastPrice'].ewm(span=p1).mean()
    # # ewm1 = m4.ewm(span=p1).mean()
    # ewm1 = m4.rolling(window=p3 * 10, min_periods=1).mean()
    # mm = m1.add(m2).add(m3).add(m4) / 4
    # dmm = mm.diff(1)
    # # print("len:data", data_len, "\nticks", len(ticks), len(price))
    # # print(len(profit), len(total_long_profit), len(mx), len(dmx))
    # # print(len(c1), len(c2), len(c3), len(c4), c1[-1:])
    # # print(profit[:5])
    # # print(dmx[-15:-5])
    # # print(ticks.head())
    # # print(type(m1),type(total_long_profit))
    # # print(ewm1[-5:-1], m4[-5:-1])
    #
    # fig = plt.figure()
    # pc = fig.add_subplot(111)
    # # pc = plt.subplot(1, 1, 1)
    # # pc.plot(mx)
    # # # pc.plot(mn)
    # # pc.plot(dmx)
    # # # pc.plot(dmn)
    # # pc.plot(profit.rolling(window=data_len, min_periods=1).sum())
    # # pc.plot(c1, color='r')
    # ra = bs[bs == 1].index
    # ga = bs[bs == -1].index
    # pc.plot(price[bs.index], color='g')
    # pc.scatter(ra, price[ra], color='r')
    # pc.scatter(ga, price[ga], color='y')
    # # pc.plot(dmm, color='r')
    # # pc.plot(sell, color='y')
    #
    # # pc.plot(dmx, color='r')
    # # pc.plot(c1, color='r')
    # # pc.plot(m3, color='y')
    # # pc.plot(c2, color='b')
    #
    # pc2 = pc.twinx()
    # pc2.plot(ama)
    # pc2.plot(mx)
    # pc2.plot(mn)
    #
    # # pc2.plot(ama)
    # # pc2.plot(mn)
    #
    # plt.legend()
    # plt.show()
    #
    # print("aaa")
